#!/usr/bin/python
# *-* coding: utf-8 *-*

import dis, opcode, imp

class SecurityException (Exception):
	def __init__ (self, message):
		Exception.__init__(self, message)
	
def getStackItemText (codeObj, item):
	if item[0] == "Name" or item[0] == "Attr":
		return codeObj.co_names[item[1]]
	elif item[0] == "Const":
		return codeObj.co_consts[item[1]]

def safetyCheck (codeObj):
	stack = []
	for i in range(len(codeObj.co_code)):
		current = ord(codeObj.co_code[i])
		if current == opcode.opmap["IMPORT_NAME"]:
			importName = codeObj.co_names[ord(codeObj.co_code[i+2]) * 255 + ord(codeObj.co_code[i+1])]
			print "One import found :", importName
			if not importName in ["base_logic", "rightsManager", "imp", "decorators", "Queue"]:
				raise SecurityException("Import of '%s' is forbidden." % importName)
		elif current == opcode.opmap["CALL_FUNCTION"]:
			argCount = ord(codeObj.co_code[i+2]) * 255 + ord(codeObj.co_code[i+1])
			print "Un appel de fonction avec %i argument(s)." % argCount
			for item in stack[-argCount:]:
				print getStackItemText(codeObj, item)
			print "Nom :", getStackItemText(codeObj, stack[-argCount-1])
			if getStackItemText(codeObj, stack[-argCount-1]) == "system":
				raise SecurityException("Execution of system is forbidden")
			stack = stack[:-argCount]
			print "---"
		elif current == opcode.opmap["LOAD_NAME"]:
			stack.append(("Name", ord(codeObj.co_code[i+2]) * 255 + ord(codeObj.co_code[i+1])))
		elif current == opcode.opmap["LOAD_CONST"]:
			stack.append(("Const", ord(codeObj.co_code[i+2]) * 255 + ord(codeObj.co_code[i+1])))
		elif current == opcode.opmap["LOAD_ATTR"]:
			stack.append(("Attr", ord(codeObj.co_code[i+2]) * 255 + ord(codeObj.co_code[i+1])))

#def safeImport (filename):
	#fh = open(filename)
	#code = fh.read()
	#fh.close()
	#codeObj = compile(code, filename, "exec")
	#dis.dis(codeObj)
	#codeController(codeObj)

#if __name__ == "__main__":
	#safeImport("demo.py")
	#if 5 == 10:
		#safeImport(safeImport(bidule.open(file("aaa"))))
